MODULE CMF_CALC_FLDSTG_MOD
!==========================================================
!* PURPOSE: calculate river and floodplain staging
!
! (C) D.Yamazaki & E. Dutra  (U-Tokyo/FCUL)  Aug 2019
!
! Licensed under the Apache License, Version 2.0 (the "License");
!   You may not use this file except in compliance with the License.
!   You may obtain a copy of the License at: http://www.apache.org/licenses/LICENSE-2.0
!
! Unless required by applicable law or agreed to in writing, software distributed under the License is 
!  distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 
! See the License for the specific language governing permissions and limitations under the License.
!==========================================================
CONTAINS
!####################################################################
! -- CMF_CALC_FLDSTG_DEF  !! default flood stage calculation
! -- CMF_OPT_FLDSTG_ES    !! optimized code for vector processor (such as Earth Simulator), activated using LSTG_ES=.TRUE. option).
! --
!####################################################################
SUBROUTINE CMF_CALC_FLDSTG_DEF
USE PARKIND1,           ONLY: JPIM, JPRB, JPRD
USE YOS_CMF_INPUT,      ONLY: NLFP
USE YOS_CMF_MAP,        ONLY: NSEQALL, D2GRAREA, D2RIVLEN, D2RIVWTH, D2RIVELV
USE YOS_CMF_MAP,        ONLY: D2RIVSTOMAX, D2FLDSTOMAX, D2FLDGRD, DFRCINC
USE YOS_CMF_PROG,       ONLY: P2RIVSTO, P2FLDSTO
USE YOS_CMF_DIAG,       ONLY: D2RIVDPH, D2FLDDPH, D2FLDFRC, D2FLDARE, D2SFCELV, D2STORGE
USE YOS_CMF_DIAG,       ONLY: P0GLBSTOPRE2, P0GLBSTONEW2, P0GLBRIVSTO, P0GLBFLDSTO, P0GLBFLDARE
IMPLICIT NONE

!*** LOCAL
INTEGER(KIND=JPIM),SAVE    :: ISEQ, I
REAL(KIND=JPRD),SAVE       :: PSTOALL
REAL(KIND=JPRB),SAVE       :: DSTOALL, DSTONOW, DSTOPRE, DWTHNOW, DWTHPRE, DDPHPRE, DWTHINC
!$OMP THREADPRIVATE        (I,DSTOALL, DSTONOW, DSTOPRE, DWTHNOW, DWTHPRE, DDPHPRE, DWTHINC,PSTOALL)
!================================================
P0GLBSTOPRE2=0._JPRD
P0GLBSTONEW2=0._JPRD
P0GLBRIVSTO =0._JPRD
P0GLBFLDSTO =0._JPRD
P0GLBFLDARE =0._JPRD

! Estimate water depth and flood extent from water storage
!   Solution for Equations (1) and (2) in [Yamazaki et al. 2011 WRR].

!$OMP PARALLEL DO REDUCTION(+:P0GLBSTOPRE2)
DO ISEQ=1, NSEQALL
!
  PSTOALL = P2RIVSTO(ISEQ,1) + P2FLDSTO(ISEQ,1)
  DSTOALL = REAL( PSTOALL, KIND=JPRB)
  P0GLBSTOPRE2     = P0GLBSTOPRE2 + PSTOALL

  IF( PSTOALL > D2RIVSTOMAX(ISEQ,1) )THEN
    I=1
    DSTOPRE = D2RIVSTOMAX(ISEQ,1)
    DWTHPRE = D2RIVWTH(ISEQ,1)
    DDPHPRE = 0._JPRB
    DWTHINC = D2GRAREA(ISEQ,1) * D2RIVLEN(ISEQ,1)**(-1._JPRB) * DFRCINC
    DO WHILE( PSTOALL > D2FLDSTOMAX(ISEQ,1,I) .AND. I<=NLFP)
      DSTOPRE = D2FLDSTOMAX(ISEQ,1,I)
      DWTHPRE = DWTHPRE + DWTHINC
      DDPHPRE = DDPHPRE + D2FLDGRD(ISEQ,1,I) * DWTHINC
      I=I+1
      IF( I>NLFP ) EXIT
    END DO
    IF( I>NLFP )THEN
      DSTONOW = DSTOALL - DSTOPRE
      DWTHNOW = 0._JPRB
      D2FLDDPH(ISEQ,1) = DDPHPRE + DSTONOW * DWTHPRE**(-1._JPRB) * D2RIVLEN(ISEQ,1)**(-1._JPRB)
    ELSE
      DSTONOW =  DSTOALL - DSTOPRE
      DWTHNOW = -DWTHPRE + &
&      ( DWTHPRE**2._JPRB + 2._JPRB * DSTONOW * D2RIVLEN(ISEQ,1)**(-1._JPRB) * D2FLDGRD(ISEQ,1,I)**(-1._JPRB) )**0.5_JPRB
      D2FLDDPH(ISEQ,1) = DDPHPRE + D2FLDGRD(ISEQ,1,I) * DWTHNOW
    ENDIF
    P2RIVSTO(ISEQ,1) = D2RIVSTOMAX(ISEQ,1) + D2RIVLEN(ISEQ,1) * D2RIVWTH(ISEQ,1) * D2FLDDPH(ISEQ,1)
    P2RIVSTO(ISEQ,1) = MIN(P2RIVSTO(ISEQ,1),PSTOALL)

    D2RIVDPH(ISEQ,1) = REAL(P2RIVSTO(ISEQ,1),JPRB) * D2RIVLEN(ISEQ,1)**(-1._JPRB) * D2RIVWTH(ISEQ,1)**(-1._JPRB)
!
    P2FLDSTO(ISEQ,1) = PSTOALL - P2RIVSTO(ISEQ,1)
    P2FLDSTO(ISEQ,1) = MAX( P2FLDSTO(ISEQ,1), 0._JPRD )
    D2FLDFRC(ISEQ,1) = (-D2RIVWTH(ISEQ,1) + DWTHPRE + DWTHNOW ) * (DWTHINC*NLFP)**(-1._JPRB)  !! bugfix 191113, (10._JPRB -> NLFP)
    D2FLDFRC(ISEQ,1) = MAX( D2FLDFRC(ISEQ,1),0._JPRB)
    D2FLDFRC(ISEQ,1) = MIN( D2FLDFRC(ISEQ,1),1._JPRB)
    D2FLDARE(ISEQ,1) = D2GRAREA(ISEQ,1)*D2FLDFRC(ISEQ,1)
  ELSE
    P2RIVSTO(ISEQ,1) = PSTOALL
    D2RIVDPH(ISEQ,1) = DSTOALL * D2RIVLEN(ISEQ,1)**(-1._JPRB) * D2RIVWTH(ISEQ,1)**(-1._JPRB)
    D2RIVDPH(ISEQ,1) = MAX( D2RIVDPH(ISEQ,1), 0._JPRB )
    P2FLDSTO(ISEQ,1) = 0._JPRD
    D2FLDDPH(ISEQ,1) = 0._JPRB
    D2FLDFRC(ISEQ,1) = 0._JPRB
    D2FLDARE(ISEQ,1) = 0._JPRB
  ENDIF
END DO
!$OMP END PARALLEL DO

!$OMP PARALLEL DO SIMD REDUCTION(+:P0GLBSTONEW2,P0GLBRIVSTO,P0GLBFLDSTO,P0GLBFLDARE)
DO ISEQ=1, NSEQALL
  D2SFCELV(ISEQ,1) = D2RIVELV(ISEQ,1) + D2RIVDPH(ISEQ,1)
  D2STORGE(ISEQ,1) = REAL(P2RIVSTO(ISEQ,1)+P2FLDSTO(ISEQ,1),KIND=JPRB)
  P0GLBSTONEW2     = P0GLBSTONEW2 + P2RIVSTO(ISEQ,1) + P2FLDSTO(ISEQ,1)
  P0GLBRIVSTO      = P0GLBRIVSTO  + P2RIVSTO(ISEQ,1)
  P0GLBFLDSTO      = P0GLBFLDSTO  + P2FLDSTO(ISEQ,1)
  P0GLBFLDARE      = P0GLBFLDARE  + D2FLDARE(ISEQ,1)
END DO
!$OMP END PARALLEL DO SIMD

END SUBROUTINE CMF_CALC_FLDSTG_DEF
!####################################################################
!
!
!
!
!####################################################################
SUBROUTINE CMF_OPT_FLDSTG_ES
! ==========
! Optional code for Earth Simulator (Vector Processor)
! Specify option: LSTG_ES=.TRUE.
! Faster computation on vector prosessor by avoiding IF-THEN function. (note this code will be slow on Scaler Processor)
! ==========
USE PARKIND1,           ONLY: JPIM, JPRB, JPRD
USE YOS_CMF_INPUT,      ONLY: NLFP
USE YOS_CMF_MAP,        ONLY: NSEQALL, D2GRAREA, D2RIVLEN, D2RIVWTH, D2RIVELV, D2RIVHGT
USE YOS_CMF_MAP,        ONLY: D2RIVSTOMAX, D2FLDSTOMAX, D2FLDGRD, DFRCINC
USE YOS_CMF_PROG,       ONLY: P2RIVSTO, P2FLDSTO
USE YOS_CMF_DIAG,       ONLY: D2RIVDPH, D2FLDDPH, D2FLDFRC, D2FLDARE, D2SFCELV, D2STORGE
USE YOS_CMF_DIAG,       ONLY: P0GLBSTOPRE2, P0GLBSTONEW2, P0GLBRIVSTO, P0GLBFLDSTO, P0GLBFLDARE
IMPLICIT NONE

!*** LOCAL
REAL(KIND=JPRB)            :: D2STODWN(NSEQALL,1)
REAL(KIND=JPRB)            :: D2WTHPRE(NSEQALL,1), D2WTHINC(NSEQALL,1)

! SAVE for OpenMP
INTEGER(KIND=JPIM),SAVE    :: ISEQ, I
REAL(KIND=JPRD),SAVE       :: PSTOALL
REAL(KIND=JPRB),SAVE       :: DSTOALL, DSTONOW, DWTHNOW
!$OMP THREADPRIVATE          (DSTOALL, DSTONOW, DWTHNOW, PSTOALL)
!================================================
P0GLBRIVSTO=0._JPRD
P0GLBFLDSTO=0._JPRD
P0GLBFLDARE=0._JPRD
P0GLBSTOPRE2=0._JPRD
P0GLBSTONEW2=0._JPRD

! [1] Assume all waters in river channel
!$OMP PARALLEL DO SIMD REDUCTION(+:P0GLBSTOPRE2)
DO ISEQ=1, NSEQALL
  PSTOALL = P2RIVSTO(ISEQ,1) + P2FLDSTO(ISEQ,1)
  DSTOALL = REAL( PSTOALL, KIND=JPRB)

  P2RIVSTO(ISEQ,1) = PSTOALL
  D2RIVDPH(ISEQ,1) = DSTOALL * D2RIVLEN(ISEQ,1)**(-1._JPRB) * D2RIVWTH(ISEQ,1)**(-1._JPRB)
  D2RIVDPH(ISEQ,1) = MAX( D2RIVDPH(ISEQ,1), 0._JPRB )
  P2FLDSTO(ISEQ,1) = 0._JPRD
  D2FLDDPH(ISEQ,1) = 0._JPRB
  D2FLDFRC(ISEQ,1) = 0._JPRB
  D2FLDARE(ISEQ,1) = 0._JPRB

  D2STODWN(ISEQ,1) = D2RIVSTOMAX(ISEQ,1)
  D2WTHPRE(ISEQ,1) = D2RIVWTH(ISEQ,1)
  D2WTHINC(ISEQ,1) = D2GRAREA(ISEQ,1) * D2RIVLEN(ISEQ,1)**(-1._JPRB) * DFRCINC

  P0GLBSTOPRE2     = P0GLBSTOPRE2 + PSTOALL
END DO
!$OMP END PARALLEL DO SIMD

! [2] Check floodplain level from I=1 to NLFP. Make I-NLFP loop outside for parallel computing (SIMD/Vector)
DO I=1, NLFP

!$OMP PARALLEL DO SIMD
  DO ISEQ=1, NSEQALL
    PSTOALL = P2RIVSTO(ISEQ,1) + P2FLDSTO(ISEQ,1)
    DSTOALL = REAL( PSTOALL, KIND=JPRB)

    DSTONOW = DSTOALL - D2STODWN(ISEQ,1)
    DSTONOW = MAX( DSTONOW, 0._JPRB )
    DWTHNOW = -D2WTHPRE(ISEQ,1) + &
&    ( D2WTHPRE(ISEQ,1)**2._JPRB + 2._JPRB * DSTONOW * D2RIVLEN(ISEQ,1)**(-1._JPRB) * D2FLDGRD(ISEQ,1,I)**(-1._JPRB) )**0.5_JPRB
    DWTHNOW = MIN( DWTHNOW, D2WTHINC(ISEQ,1) )
    DWTHNOW = MAX( DWTHNOW, 0._JPRB )              !! modify v4.04

    D2FLDDPH(ISEQ,1) = D2FLDDPH(ISEQ,1) + D2FLDGRD(ISEQ,1,I) * DWTHNOW
    D2FLDFRC(ISEQ,1) = D2FLDFRC(ISEQ,1) + DWTHNOW/D2WTHINC(ISEQ,1) * NLFP**(-1._JPRB)

    !! Update downside floodplain step storage/depth/width 
    D2STODWN(ISEQ,1) = D2FLDSTOMAX(ISEQ,1,I)
    D2WTHPRE(ISEQ,1) = D2WTHPRE(ISEQ,1) + D2WTHINC(ISEQ,1)
  END DO
  !$OMP END PARALLEL DO SIMD
END DO

!! [3] flood extent saturated case
!$OMP PARALLEL DO SIMD
DO ISEQ=1, NSEQALL
  PSTOALL = P2RIVSTO(ISEQ,1) + P2FLDSTO(ISEQ,1)
  DSTOALL = REAL( PSTOALL, KIND=JPRB)
  DSTONOW = DSTOALL - D2STODWN(ISEQ,1)
  DSTONOW = MAX( DSTONOW, 0._JPRB )
  D2FLDDPH(ISEQ,1) = D2FLDDPH(ISEQ,1) + DSTONOW * D2WTHPRE(ISEQ,1)**(-1._JPRB) * D2RIVLEN(ISEQ,1)**(-1._JPRB)
END DO
!$OMP END PARALLEL DO SIMD

!! [4] Floodplain stage diagnose
!$OMP PARALLEL DO SIMD
DO ISEQ=1, NSEQALL
!  IF( D2FLDDPH(ISEQ,1)>0 )THEN !! bugfix v4.04
  IF( D2FLDDPH(ISEQ,1) > (10._JPRB**-5) )THEN !! bugfix v4.04, to avoid false positive FLDDPH due to rounding error.
    PSTOALL = P2RIVSTO(ISEQ,1) + P2FLDSTO(ISEQ,1)

    D2RIVDPH(ISEQ,1) = D2RIVHGT(ISEQ,1) + D2FLDDPH(ISEQ,1)
    P2RIVSTO(ISEQ,1) = D2RIVLEN(ISEQ,1) * D2RIVWTH(ISEQ,1) * D2RIVDPH(ISEQ,1)
    P2RIVSTO(ISEQ,1) = MIN( P2RIVSTO(ISEQ,1), PSTOALL )           !! modify v4.04
!
    P2FLDSTO(ISEQ,1) = PSTOALL - P2RIVSTO(ISEQ,1)
    P2FLDSTO(ISEQ,1) = MAX( P2FLDSTO(ISEQ,1), 0._JPRD )

    D2FLDFRC(ISEQ,1) = MAX( D2FLDFRC(ISEQ,1),0._JPRB)
    D2FLDFRC(ISEQ,1) = MIN( D2FLDFRC(ISEQ,1),1._JPRB)
    D2FLDARE(ISEQ,1) = D2GRAREA(ISEQ,1)*D2FLDFRC(ISEQ,1)
  ENDIF
END DO
!$OMP END PARALLEL DO SIMD

!$OMP PARALLEL DO SIMD REDUCTION(+:P0GLBSTONEW2,P0GLBRIVSTO,P0GLBFLDSTO,P0GLBFLDARE)
DO ISEQ=1, NSEQALL
  D2SFCELV(ISEQ,1) = D2RIVELV(ISEQ,1) + D2RIVDPH(ISEQ,1)
  D2STORGE(ISEQ,1) =REAL(P2RIVSTO(ISEQ,1)+P2FLDSTO(ISEQ,1),KIND=JPRB)
  P0GLBSTONEW2      = P0GLBSTONEW2+ P2RIVSTO(ISEQ,1) + P2FLDSTO(ISEQ,1)
  P0GLBRIVSTO       = P0GLBRIVSTO + P2RIVSTO(ISEQ,1)
  P0GLBFLDSTO       = P0GLBFLDSTO + P2FLDSTO(ISEQ,1)
  P0GLBFLDARE       = P0GLBFLDARE + D2FLDARE(ISEQ,1)
END DO
!$OMP END PARALLEL DO SIMD

END SUBROUTINE CMF_OPT_FLDSTG_ES
!####################################################################
END MODULE CMF_CALC_FLDSTG_MOD
